open Facile
open Easy

open Uav_data
open Target

let pa = fun a -> 
  Array.iter (fun t -> Printf.printf "%2d " (Fd.elt_value t)) a; 
  print_newline();
  flush stdout
     

let position_g = fun task pos card ->
  Goals.create
    (fun () ->
     (* position in plans *)
     (* constraints on position *)
     let vars_pos = 
       Array.init m 
	 (fun i ->
	    let l =ref [] in
	      Array.iteri 
		(fun k t ->
		   let t = Fd.elt_value t in
		     if t=i then l:= pos.(k)::!l) task;
	      Array.of_list !l) in

       for i=0 to m-1 do
        Cstr.post (Alldiff.cstr vars_pos.(i));
        Cstr.post (fd2e pos.(n+i) =~ i2e 0);
        Cstr.post (fd2e pos.(n+m+i) =~ i2e (Fd.elt_value card.(i) -1))
       done;

	 (Goals.Array.forall Goals.indomain pos)
  )


let load = Array.init m (fun i -> Invariant.create 0)

let print_load = fun () ->
  Printf.printf "Loads: ";
  Array.iter (fun a -> Printf.printf "%d %!" (Invariant.get a)) load;
  print_newline()

let maintaining_load_cstr task = 
  let name = "maintaining_load"
  and delay ct = 
    Array.iteri
      (fun i v -> Fd.delay [Var.Fd.on_subst] v ~waking_id:i ct) task
  and update id = 
    let k = Fd.elt_value task.(id) in
      Invariant.set load.(k) (Invariant.get load.(k) +1);
      (*print_load ();*)
      false in
    Cstr.create ~name ~nb_wakings:(n+2*m) update delay 


let insertion_cost = fun k t p i ->
  let p = Array.of_list p in
    time.(k).(p.(i-1)).(t) + time.(k).(t).(p.(i)) - time.(k).(p.(i-1)).(p.(i))

let rec insertion = fun t p i ->
  if i = 0
  then t :: p
  else List.hd p :: insertion t (List.tl p) (i-1)

(* use the insertion heuristic to schedule tasks in plans *)
let glutony_schedule = fun task pos () -> 
  (* we define the tasks that each uav has to do *)
  let uav_tasks = Array.init m (fun _ -> []) in
  let plans = Array.init m (fun i -> [n+i; n+m+i]) in
    for i=0 to n-1 do
      let t = Fd.elt_value task.(i) in
      uav_tasks.(t) <- i :: uav_tasks.(t)
    done;

    (* for each aircraft *)
    for k=0 to m-1 do
      while List.length uav_tasks.(k) > 0 do
        let t = List.hd uav_tasks.(k)
        and position = ref (-1)
        and lower_cost = ref max_int
      in
      uav_tasks.(k) <- List.tl uav_tasks.(k);

	    for i=1 to List.length plans.(k) -1 do
	      let ic = insertion_cost k t plans.(k) i in
        if ic < !lower_cost then
          begin
            lower_cost := ic;
            position := i
          end
      done;
      plans.(k) <- insertion t plans.(k) !position;
    done;
    Array.iteri (fun i t -> ignore (Fd.unify pos.(t) i)) 
    (Array.of_list plans.(k));
  done

let t_ref = Sys.time ()  


let _ =
  let task = Array.init (n+2*m) (fun _ -> Fd.interval 0 (m-1)) in


  let next = Array.init (n+2*m) (fun _ -> Fd.interval 0 (n+2*m-1)) in
  let prev = Array.init (n+2*m) (fun _ -> Fd.interval 0 (n+2*m-1)) in

  let pos = Array.init (n+2*m) (fun _ -> Fd.interval 0 (n+2*m-1)) in

  let cost = Array.init m 
    (fun i ->
      let c = Array.init (n+2*m) 
        (fun k -> 
      fd2e
        (FdArray.get 
          (Array.map (fun x -> Fd.interval x x) time.(i).(k))
          next.(k)) *~ (fd2e (Reify.boolean  (fd2e task.(k) =~ i2e i )))
	    
        ) in
        Arith.e2fd (Arith.sum c)
    ) in

  let final_cost = FdArray.max cost in
  
  let card = (*Array.init m (fun _ -> Fd.interval 0 (Array.length task))*)
    
    Array.init m 
      (fun i ->
        Arith.e2fd 
          (Arith.sum_fd 
            (Array.map 
              (fun task_k -> Reify.boolean (fd2e task_k =~ i2e i)) task)
          )
      ) in

  (***************)
  (* Constraints *)
  (***************)

(* *)
  (*Cstr.post (Card.cstr card task);*)

  (* maintain load for the search heuristic *)
  Cstr.post (maintaining_load_cstr task);

  (* next and prev definition constraints *)
  Cstr.post (Alldiff.cstr next);
  Cstr.post (Alldiff.cstr prev);
  for k=0 to m-1 do
    Cstr.post (fd2e prev.(n+k) =~ i2e (n+m+k));
    Cstr.post (fd2e next.(n+m+k) =~  i2e (n+k));
  done;
  for k=0 to (n+2*m-1) do
    Cstr.post (fd2e next.(k) <>~ i2e k);
    for p= 0 to (n+2*m-1) do
      if k<>p then begin
        Cstr.post ((fd2e next.(k) =~ i2e p) <=>~~ (fd2e prev.(p) =~ i2e k));
        Cstr.post ((fd2e next.(k) =~ i2e p) =>~~ 
            (fd2e task.(k) =~ fd2e task.(p)));
        Cstr.post ((fd2e prev.(k) =~ i2e p) =>~~ 
            (fd2e task.(k) =~ fd2e task.(p)));
      end
    done
  done;

  (* instantiate takeoff and landing tasks *)
  for k=0 to m-1 do
    Cstr.post (fd2e task.(n+k) =~ i2e k);
    Cstr.post (fd2e task.(n+m+k) =~ i2e k);
    (*Cstr.post (fd2e pos.(n+k) =~ i2e 0);
    Cstr.post (fd2e pos.(n+m+k) =~ fd2e card.(k) -~ i2e 1);*)
  done;

  (* video capacity constraints *)
  for i=0 to m-1 do
    for k=0 to n-1 do
      if (ttype.(k) = Se || ttype.(k) = Id) && (video.(i) = 0)
      then 
        Cstr.post (fd2e task.(k) <>~ i2e i);
    done
  done;

  (* drop capacity constraints *)
  let drop_tasks_index = (* compute the index of the drop tasks *)
    let index = Array.init n (fun k -> k) in
      Array.of_list 
        (List.filter (fun x-> ttype.(x) = Dr) 
    (Array.to_list index))
  in
  let drops_in_plan = Array.init m 
    (fun i ->
      Arith.e2fd
        (Arith.sum_fd
    (Array.map 
      (fun tid -> Reify.boolean 
          (fd2e task.(tid) =~ i2e i)) drop_tasks_index)))
  in
  for i=0 to m-1 do
      Cstr.post (fd2e drops_in_plan.(i) <=~ i2e ball.(i))
  done;

  (* card posistion constraints *)
  for i=0 to n+2*m-1 do
    Cstr.post (fd2e pos.(i) <=~ fd2e (FdArray.get card task.(i)) -~ i2e 1);
  done;

  (* position and next constraints *)
  for i1 = 0 to n-1 do
    for i2 = 0 to n-1 do
      if i1<>i2 then
        Cstr.post (((fd2e pos.(i1) =~ fd2e pos.(i2) +~ i2e 1) &&~~
          (fd2e task.(i1) =~ fd2e task.(i2))) <=>~~
          (fd2e next.(i1) =~ i2e i2));
    done
  done;
  
  
  
  (* nicolas' vrp constraint *)
  Cstr.post (Vrp.cstr n m prev next);
  
  let goal =
  
    let min_load_labeling = 
      Goals.Array.forall
        (Goals.instantiate 
          (Domain.choose 
          (fun a1 a2 -> Invariant.get load.(a1) < Invariant.get load.(a2))
          ))
    in
    let random_labeling = 
      Goals.Array.forall
        (Goals.instantiate
          (Domain.choose
            (fun a1 a2 -> Random.bool ())))
    in
	 
    let labeling = ref min_load_labeling in
    let schedule = random_labeling prev in
    let g_schedule = Goals.atomic (glutony_schedule task pos) in
    let sched_t = ref true in

    Arg.parse [ 
      ("-simple_labeling", 
      Arg.Unit (fun () -> labeling := Goals.Array.labeling), 
      "Use simple labeling strategy for allocation"); 
      ("-random_labeling",
      Arg.Unit (fun () -> labeling := random_labeling),
      "Use random labeling");
      ("-glutony",
      Arg.Clear sched_t,
      "Use glutony insertion to schedule allocation")
    ] (fun _ -> ()) "Usage: ";

    if !sched_t 
    then (!labeling task) &&~ (position_g task pos card ()) &&~ (schedule)
    else (!labeling task) &&~ g_schedule &&~ (Goals.Array.forall Goals.indomain cost)
      (*(!labeling task) &&~ (schedule)*)
  in

let print_costs = fun () ->
  Printf.printf "% 4.2f\t" (Sys.time () -. t_ref); 
  Fd.fprint stdout final_cost;
  Printf.printf "\t\t\t";
  Array.iter (fun c -> Fd.fprint stdout c; Printf.printf "\t") cost;
  print_newline() in



let opt_goal = Goals.minimize ~step: 1 goal final_cost (fun _ -> print_costs ())
in

  ignore(Goals.solve ((opt_goal &&~ Goals.fail) ||~ Goals.success));
  Printf.printf "End of search: % 4.2f\n%!" (Sys.time() -. t_ref)
